# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 10:57 
@Author  : ZhangShenao 
@File    : 8.基于Token长度的分割器.py 
@Desc    : 基于Token长度的分割器

使用tiktoken包,可以相对精确地计算token数量
基于此分割的文本长度会更加精确
"""
import tiktoken
from langchain_community.document_loaders import UnstructuredFileLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter


def _calc_token_count(text: str) -> int:
    """
    计算文本所包含的token数量
    :param text: 文本
    :return: token数量
    """

    # 基于Embedding模型,创建token编码器
    encoder = tiktoken.encoding_for_model('text-embedding-3-small')

    # 对文本进行token编码,生成token列表,返回列表长度即为token数量
    return len(encoder.encode(text))


# 加载文档
loader = UnstructuredFileLoader('./docs/骆驼祥子.txt')
docs = loader.load()

# 创建RecursiveCharacterTextSplitter
separators = [
    '\n\n',
    '\n',
    r'。|！|？',
    r'\.\s|\!\s|\?\s',  # 英文标点符号后面通常需要加空格
    r'；|;\s',
    r'，|,\s',
    ' ',
    ''
]
splitter = RecursiveCharacterTextSplitter(
    separators=separators,  # 指定自定义分隔符列表
    chunk_size=20,  # 指定文档分块大小
    chunk_overlap=5,  # 指定文档块之间的重叠部分大小
    add_start_index=True,  # 添加文档块的起始索引
    is_separator_regex=True,  # 指定分隔符为正则表达式形式
    length_function=_calc_token_count  # 指定计算token数量的函数
)

# 等价于直接调用类方法from_tiktoken_encoder()来快速创建基于tiktoken分词器的文本分割器
RecursiveCharacterTextSplitter.from_tiktoken_encoder(model_name='text-embedding-3-small')

# 分割文档
chunks = splitter.split_documents(docs)
for chunk in chunks:
    print(f'分片长度: {len(chunk.page_content)}, 文本内容: {chunk.page_content}, 元数据: {chunk.metadata}')
